"""
资源监控模块 - 负责获取CPU、内存和网络数据
"""

import psutil
import time

# 上一次网络计数和时间戳
last_net_io = None
last_time = None

def get_cpu_usage():
    """获取CPU使用率"""
    return psutil.cpu_percent(interval=0)

def get_memory_usage():
    """获取内存使用率"""
    return psutil.virtual_memory().percent

def format_speed(bytes_per_sec):
    """将字节速率转换为可读格式 (KB/s 或 MB/s)"""
    if bytes_per_sec < 1024 * 1024:  # 小于1MB/s
        return f"{bytes_per_sec / 1024:.1f} KB/s"
    else:  # 大于等于1MB/s
        return f"{bytes_per_sec / (1024 * 1024):.1f} MB/s"

def get_network_speed():
    """获取当前网络速度（上传和下载）
    
    返回:
        tuple: (上传速度, 下载速度) 字符串格式 (例如 "1.2 KB/s")
        首次调用返回 ("--", "--")
    """
    global last_net_io, last_time
    
    # 获取当前网络IO统计和时间戳
    current_net_io = psutil.net_io_counters()
    current_time = time.time()
    
    # 首次调用时没有上一次的数据，无法计算速度
    if last_net_io is None or last_time is None:
        last_net_io = current_net_io
        last_time = current_time
        return "--", "--"  # 首次调用返回占位符
    
    # 计算时间差（秒）
    time_diff = current_time - last_time
    
    # 确保时间差大于0以避免除零错误
    if time_diff <= 0:
        return "--", "--"
    
    # 计算上传和下载的字节差
    upload_bytes = current_net_io.bytes_sent - last_net_io.bytes_sent
    download_bytes = current_net_io.bytes_recv - last_net_io.bytes_recv
    
    # 计算速度（字节/秒）
    upload_speed = upload_bytes / time_diff
    download_speed = download_bytes / time_diff
    
    # 更新上一次数据
    last_net_io = current_net_io
    last_time = current_time
    
    # 格式化并返回速度
    return format_speed(upload_speed), format_speed(download_speed)

# 简单测试代码
if __name__ == "__main__":
    print("资源监控测试:")
    print(f"CPU 使用率: {get_cpu_usage()}%")
    print(f"内存使用率: {get_memory_usage()}%")
    print("等待2秒收集网络数据...")
    get_network_speed()  # 首次调用初始化基准数据
    time.sleep(2)
    upload, download = get_network_speed()
    print(f"上传速度: {upload}")
    print(f"下载速度: {download}") 